import os
import cv2
import argparse
from tqdm import tqdm


def extract_frames(video_path, output_dir, interval=3, ext='jpg'):
    """
    从单个视频中抽帧保存到对应文件夹
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频: {video_path}")
        return

    os.makedirs(output_dir, exist_ok=True)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    video_name = os.path.splitext(os.path.basename(video_path))[0]
    saved = 0

    for i in tqdm(range(frame_count), desc=f"Extracting {video_name}"):
        ret, frame = cap.read()
        if not ret:
            break
        if i % interval == 0:
            save_path = os.path.join(output_dir, f"{i:06d}.{ext}")
            cv2.imwrite(save_path, frame)
            saved += 1

    cap.release()
    print(f"{video_name}: 共抽取 {saved} 帧 (间隔 {interval})")


def main(root_dir, output_root, interval=3, ext='jpg'):
    """
    遍历视频目录 video/，每个视频一个输出子文件夹 frames/视频名/
    """
    os.makedirs(output_root, exist_ok=True)
    videos = [f for f in os.listdir(root_dir)
              if f.lower().endswith(('.mp4', '.avi', '.mov', '.mkv'))]

    if not videos:
        print(f"未在 {root_dir} 下找到任何视频文件")
        return

    for v in videos:
        video_path = os.path.join(root_dir, v)
        video_name = os.path.splitext(v)[0]
        output_dir = os.path.join(output_root, video_name)
        if os.path.exists(output_dir) and len(os.listdir(output_dir)) > 0:
            print(f"跳过已抽帧的视频: {v}")
            continue
        extract_frames(video_path, output_dir, interval=interval, ext=ext)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="简单视频抽帧脚本")
    parser.add_argument("--root_dir", type=str,
                        default=r"F:\project\jili_three_project\foggy_project_all\1108_new\video",
                        help="输入视频目录 (里面全是视频)")
    parser.add_argument("--output_root", type=str,
                        default=r"F:\project\jili_three_project\foggy_project_all\1108_new\frames",
                        help="输出图片目录")
    parser.add_argument("--interval", type=int, default=5,
                        help="每隔多少帧抽取一帧 (默认5)")
    parser.add_argument("--ext", type=str, default="jpg",
                        choices=["jpg", "png", "webp"], help="图片格式")

    args = parser.parse_args()
    main(args.root_dir, args.output_root, args.interval, args.ext)
